from typing import Optional

from PIL import Image
import pdfplumber

from ..book import Book, Page, Content, ContentType, TableContent
from .exceptions import PageOutOfRangeException
from ..book.content import ImageContent
from ..utils import LOG


class PDFParser:
    def __init__(self):
        pass

    def paser_pdf(self, pdf_file_path: str, pages: Optional[int] = None) -> Book:
        book = Book(pdf_file_path)
        with pdfplumber.open(pdf_file_path) as pdf:
            if pages is not None and pages > len(pdf.pages):
                raise PageOutOfRangeException(len(pdf.pages), pages)

            if pages is None:
                pages_to_parse = pdf.pages
            else:
                pages_to_parse = pdf.pages[:pages]

            for page_pdf in pages_to_parse:
                page = Page(page_pdf.layout)
                raw_text = page_pdf.extract_text()
                tables = page_pdf.extract_tables()
                images = page_pdf.images
                # 提取的文本替换掉表格的
                for table_data in tables:
                    for row in table_data:
                        for cell in row:
                            raw_text = raw_text.replace(cell, '', 1)

                # text
                if raw_text:
                    raw_text_lines = raw_text.splitlines()
                    cleaned_raw_text_lines = [line.strip() for line in raw_text_lines if line.strip()]
                    cleaned_raw_text = "\n".join(cleaned_raw_text_lines)

                    text_content = Content(content_type=ContentType.TEXT, original=cleaned_raw_text)

                    page.add_content(text_content)
                    LOG.debug(f"[raw_text]\n {cleaned_raw_text}")
                # Image
                # if images:
                    # images = ImageContent(images)
                    # page.add_content(images)
                    # for s_image in images:
                    #     image = Image.open(s_image['stream'].get_data())
                    #     image.show()
                # table
                if tables:
                    table = TableContent(tables)
                    page.add_content(table)
                    LOG.debug(f"[table]\n {table}")

                book.add_page(page)

        return book
